#include "inc/core/mover1d.h"
#include "inc/core/lib.h"
#include "stdio.h"
#include "stdlib.h"
#include "math.h"


/*INTERNAL FUNCTIONS*/

__device__ floatdig mirror_b_field(timevar t,  floatdig pos)
{
	return B0;
}

__device__ floatdig uniform_e_field(timevar t, floatdig pos)
{
	floatdig Dz = pos - PULSECENTRE_ZPOS;
	floatdig DtI = t - (PULSEDELAY);
	floatdig DtD = t - (PULSEDELAY + PULSEINCTIME);
	floatdig Az = exp(- pow((Dz / PULSEWIDTH / 0.3),6));
	floatdig AtI = pow(DtI / PULSEINCTIME / 0.1, 2);
	floatdig AtD = pow(DtD / PULSEDELTIME / 0.4, 2);
	floatdig osc = sin(WCE * t);
	if ((DtI >= 0) && (DtI <= PULSEINCTIME)) 
	{
		return E0 * Az * (1 - exp(-AtI)) * osc;
	} 
	else 
	{
		if ((DtD >= 0) && (DtD <= PULSEDURATION)) 
		{
			return E0 * Az* (1/(1+AtD)) * osc;
		} 

			return 0;
	}
}

//use numofthreads == arraysize
//zpool is the array we sort
//vxpool and vypool are reaarrangeg according to changes in zpool
//thanks to edeldm from habr.ru
__global__ void qsortpools(floatdig *zpool, floatdig* vxpool, floatdig* vypool, unsigned long size)
{
	unsigned long tid = blockIdx.x*blockDim.x + threadIdx.x;
	unsigned long  k_tid = 0;
	floatdig z_tid = zpool[tid];
	floatdig vx = vxpool[tid];
	floatdig vy = vypool[tid];


	for (unsigned long i = 0; i<size; ++i)
		if (zpool[i]<z_tid)
			++k_tid;

	__syncthreads();

	zpool[k_tid] = z_tid;
	vxpool[k_tid] = vx;
	vypool[k_tid] = vx;
}

__global__ void update_external_fields(floatdig* E, floatdig* B, floatdig* zpool, timevar t)
{
	unsigned long tid = blockIdx.x*blockDim.x + threadIdx.x;
	E[tid] = uniform_e_field(t, zpool[tid]);
	B[tid] = mirror_b_field(t, zpool[tid ]);
}

__device__ void meanspeed2(speed3d* olds, speed3d* news, speed3d* mean)
{
	mean->v[VX] = (fabs(olds->v[VX])*fabs(news->v[VX]));
	mean->v[VY] = (fabs(olds->v[VY])*fabs(news->v[VY]));
	mean->v[VZ] = (fabs(olds->v[VZ])*fabs(news->v[VZ]));
}

__host__ void copypoolHostToDevice(floatdig* vxpool, floatdig* vypool, floatdig* vzpool, floatdig* zpool,
                                   particle* particles, int shift)
{
	particle* p = particles;
	floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;

	if(zpool != NULL) {
		for (i = 0; i < PARTICLEPOOL; i++) {
			temp[i] = p[i + shift].pos.v[Z];
		}
		CUDA_CALL(cudaMemcpy(zpool, temp,  sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	}

	if(vxpool != NULL) {
		for (i = 0; i < PARTICLEPOOL; i++) {
			temp[i] = p[i + shift].speed.v[VX];
		}
		CUDA_CALL(cudaMemcpy(vxpool, temp,  sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	}

	if(vypool != NULL) {
		for (i = 0; i < PARTICLEPOOL; i++) {
			temp[i] = p[i + shift].speed.v[VY];
		}
		CUDA_CALL(cudaMemcpy(vypool, temp,  sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	}

	if(vzpool != NULL) {
		for (i = 0; i < PARTICLEPOOL; i++) {
			temp[i] = p[i + shift].speed.v[VZ];
		}
		CUDA_CALL(cudaMemcpy(vzpool, temp,  sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyHostToDevice));
	}

	free(temp);
}

__host__ void copypoolDeviceToHost(floatdig* vxpool, floatdig* vypool, floatdig* vzpool,
                                   floatdig* mvxpool, floatdig* mvypool, floatdig* zpool,
                                   particle* particles, int shift)
{
	particle* p = particles;
	floatdig* temp = (floatdig*)malloc(sizeof(floatdig)*PARTICLEPOOL);
	int i;

	if(zpool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, zpool,  sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].pos.v[Z] = temp[i];
		}
	}

	if(vxpool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, vxpool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].speed.v[VX] = temp[i];
		}
	}

	if (vypool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, vypool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].speed.v[VY] = temp[i];
		}
	}

	if (mvxpool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, mvxpool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].meanspeed2.v[VX] = temp[i];
		}
	}

	if (mvypool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, mvypool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].meanspeed2.v[VY] = temp[i];
		}
	}

	if (vzpool != NULL) {
		CUDA_CALL(cudaMemcpy(temp, vzpool, sizeof(floatdig)*PARTICLEPOOL, cudaMemcpyDeviceToHost));
		for (i = 0; i < PARTICLEPOOL; i++) {
			p[i + shift].speed.v[VZ] = temp[i];
		}
	}

	free(temp);
}

__global__ void do_backtracing(floatdig* vxpool, floatdig* vypool,
                               floatdig* E, floatdig* B, floatdig delta, floatdig qm, floatdig qmc)
{
	unsigned long nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
	floatdig A = qmc*B[nparticle]*delta/2;
	floatdig D = 1 + A*A;;
	floatdig C1 = vxpool[nparticle] + vypool[nparticle]*A + qm*delta*E[nparticle];
	floatdig C2 = vypool[nparticle] - vxpool[nparticle]*A;
	vxpool[nparticle] = (C1 + A*C2)/D;
	vypool[nparticle] = (C2 - A*C1)/D;
}

__global__ void do_advance(floatdig* vxpool, floatdig* vypool, floatdig* vzpool,
                           floatdig* mvxpool, floatdig* mvypool, floatdig* zpool,
                           floatdig* E, floatdig* B, floatdig delta, floatdig qm, floatdig qmc)
{
	unsigned long nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
	floatdig A = qmc*B[nparticle]*delta/2;
	floatdig D = 1 + A*A;;
	floatdig C1 = vxpool[nparticle] + vypool[nparticle]*A + qm*delta*E[nparticle];
	floatdig C2 = vypool[nparticle] - vxpool[nparticle]*A;
	mvxpool[nparticle] = fabs(vxpool[nparticle]);
	mvypool[nparticle] = fabs(vypool[nparticle]);
	vxpool[nparticle] = (C1 + A*C2)/D;
	vypool[nparticle] = (C2 - A*C1)/D;
	mvxpool[nparticle] *= fabs(vxpool[nparticle]);
	mvypool[nparticle] *= fabs(vypool[nparticle]);
	zpool[nparticle] += vzpool[nparticle]*delta;
}



/*API*/
mover1d* mover1d_setup_mover()
{
	floatdig grid_speed = 0;
	mover1d* mover = NULL;
	int cycles = 0;
	floatdig* B;
	floatdig* E;
	floatdig* a;
	floatdig* vxpool = NULL;
	floatdig* vypool = NULL;
	floatdig* mvxpool = NULL;
	floatdig* mvypool = NULL;
	floatdig* vzpool = NULL;
	floatdig* zpool = NULL;
	//int* flags = NULL;
	int i;
#ifdef VERBOSE
	printf("mover1d_setup_mover: initialization\n");
#endif

	CUDA_CALL(cudaMalloc(&mover, sizeof(mover1d)));
	CUDA_CALL(cudaMalloc(&B, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&E, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&zpool, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&vxpool, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&vypool, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&vzpool, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&mvxpool, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMalloc(&mvypool, PARTICLEPOOL*sizeof(floatdig)));
	//CUDA_CALL(cudaMalloc(&flags, num_of_cells*sizeof(int)));

	CUDA_CALL(cudaMemcpy(&mover->vxpool, &vxpool, sizeof vxpool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->vypool, &vypool, sizeof vypool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->vzpool, &vzpool, sizeof vypool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->mvxpool, &mvxpool, sizeof mvxpool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->mvypool, &mvypool, sizeof mvypool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->zpool, &zpool, sizeof zpool, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->E, &E, sizeof E, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemcpy(&mover->B, &B, sizeof B, cudaMemcpyHostToDevice));
	//CUDA_CALL(cudaMemcpy(&mover->flags, &flags, sizeof flags, cudaMemcpyHostToDevice));


	CUDA_CALL(cudaMemcpy(&(mover->cycles), &cycles, sizeof cycles, cudaMemcpyHostToDevice));
	CUDA_CALL(cudaMemset(E, 0,  PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(B, 0,  PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(zpool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(vxpool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(vypool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(vzpool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(mvxpool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(mvypool, 0, PARTICLEPOOL*sizeof(floatdig)));
	CUDA_CALL(cudaMemset(zpool, 0, PARTICLEPOOL*sizeof(floatdig)));
	//CUDA_CALL(cudaMemset(flags, 0, PARTICLEPOOL*sizeof(int)));
#ifdef VERBOSE
	printf("mover1d_setup_mover: initilization is succesfull\n");
#endif
	return mover;

}

void mover1d_clear_mover(mover1d* mover)
{
	floatdig* B;
	floatdig* E;
	floatdig* zpool;
	floatdig* vxpool;
	floatdig* vypool;
	floatdig* vzpool;
	floatdig* mvxpool;
	floatdig* mvypool;

#ifdef VERBOSE
	printf("mover1d_clear_mover: cleanup mover\n");
#endif

	CUDA_CALL(cudaMemcpy(&B, &(mover->B), sizeof B, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&E, &(mover->E), sizeof E, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&zpool, &(mover->zpool), sizeof zpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vxpool, &(mover->vxpool), sizeof vxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vypool, &(mover->vypool), sizeof vypool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vzpool, &(mover->vzpool), sizeof vzpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&mvxpool, &(mover->mvxpool), sizeof mvxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&mvypool, &(mover->mvypool), sizeof mvypool, cudaMemcpyDeviceToHost));

	CUDA_CALL(cudaFree(B));
	CUDA_CALL(cudaFree(E));
	CUDA_CALL(cudaFree(vxpool));
	CUDA_CALL(cudaFree(vypool));
	CUDA_CALL(cudaFree(vzpool));
	CUDA_CALL(cudaFree(mvxpool));
	CUDA_CALL(cudaFree(mvypool));
	CUDA_CALL(cudaFree(zpool));
	CUDA_CALL(cudaFree(mover));
#ifdef VERBOSE
	printf("mover1d_clear_mover: cleanup mover is succesfull\n");
#endif
}

__host__ void mover1d_backtracing(mover1d* mover,  particles_set* set, timevar dt)
{
	int nbunches = set->num_of_particles/PARTICLEPOOL;
	floatdig* B;
	floatdig* E;
	floatdig* vxpool;
	floatdig* vypool;
	floatdig* zpool;
	floatdig qm;
	floatdig qmc;
	int i;

#ifdef DEBUG
	printf("mover1d_backtracing: start\n");
#endif

	CUDA_CALL(cudaMemcpy(&B, &(mover->B), sizeof B, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&E, &(mover->E), sizeof E, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vxpool, &(mover->vxpool), sizeof vxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vypool, &(mover->vypool), sizeof vypool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&zpool, &(mover->zpool), sizeof zpool, cudaMemcpyDeviceToHost));

	for (i = 0; i < nbunches; i++) {
		copypoolHostToDevice(vxpool, vypool, NULL, zpool, set->particles, i*PARTICLEPOOL);
		update_external_fields<<<NUMOFTHREADBLOCKS,NUMOFTHREADSPERBLOCK>>>(E, B, zpool, 0);
		do_backtracing<<<NUMOFTHREADBLOCKS,NUMOFTHREADSPERBLOCK>>>(vxpool, vypool, E, B,  dt, QM, QMC);
		copypoolDeviceToHost(vxpool, vypool, NULL, NULL, NULL, zpool, set->particles, i*PARTICLEPOOL);
	}
#ifdef DEBUG
	printf("mover1d_backtracing: done\n");
#endif
}

__host__ void mover1d_advance(mover1d* mover, particles_set* set, timevar dt)
{
	int nbunches = set->num_of_particles/PARTICLEPOOL;
	floatdig* B;
	floatdig* E;
	floatdig* zpool;
	floatdig* vxpool;
	floatdig* vypool;
	floatdig* vzpool;
	floatdig* mvxpool;
	floatdig* mvypool;
	floatdig qm;
	floatdig qmc;
	int i;
	int cycles;

#ifdef DEBUG
	printf("mover1d_advance: start\n");
#endif

	CUDA_CALL(cudaMemcpy(&B, &(mover->B), sizeof B, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&E, &(mover->E), sizeof E, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&zpool, &(mover->zpool), sizeof zpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vxpool, &(mover->vxpool), sizeof vxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vypool, &(mover->vypool), sizeof vypool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&vzpool, &(mover->vzpool), sizeof vzpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&mvxpool, &(mover->mvxpool), sizeof mvxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&mvypool, &(mover->mvypool), sizeof mvypool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&cycles, &(mover->cycles), sizeof cycles, cudaMemcpyDeviceToHost));

#ifdef DEBUG
	printf("mover1d_advance: cycle %d\n", cycles);
	printf("ct is %e\n", cycles*dt);
#endif

	for (i = 0; i < nbunches; i++) {
		copypoolHostToDevice(vxpool, vypool, vzpool, zpool, set->particles, i*PARTICLEPOOL);
		update_external_fields<<<NUMOFTHREADBLOCKS,NUMOFTHREADSPERBLOCK>>>(E, B, zpool, cycles*dt);
		do_advance<<<NUMOFTHREADBLOCKS,NUMOFTHREADSPERBLOCK>>>(vxpool, vypool, vzpool, mvxpool, mvypool, zpool, E, B,  dt, QM, QMC);
		copypoolDeviceToHost(vxpool, vypool, vzpool, mvxpool, mvypool, zpool, set->particles, i*PARTICLEPOOL);
	}
	cycles++;
	CUDA_CALL(cudaMemcpy(&(mover->cycles), &cycles, sizeof cycles, cudaMemcpyHostToDevice));
#ifdef DEBUG
	printf("mover1d_advance: done\n");
#endif
}

__global__ void do_postproceesing(floatdig* d_vxpool, floatdig* d_vypool, floatdig* d_zpool, floatdig* d_rvx, floatdig* d_rvy, floatdig minz, floatdig maxz)
{
	unsigned long nparticle = GETPROCPARTICLE(blockDim, blockIdx, threadIdx);
	if(d_zpool[nparticle] < minz) {
		d_zpool[nparticle] = maxz;
		d_vxpool[nparticle] = d_rvx[nparticle];
		d_vypool[nparticle] = d_rvy[nparticle];
	}
	if(d_zpool[nparticle] > maxz) {
		d_zpool[nparticle] = minz;
		d_vxpool[nparticle] = d_rvx[nparticle];
		d_vypool[nparticle] = d_rvy[nparticle];
	}
}

__host__ void mover1d_postprocessing(mover1d* mover, particles_set* set, floatdig minz, floatdig maxz, floatdig sigma)
{
	floatdig* d_zpool;
	floatdig* d_vxpool;
	floatdig* d_vypool;
	floatdig* vzpool;
	floatdig* d_rvx;
	floatdig* d_rvy;
	int nbunches = set->num_of_particles/PARTICLEPOOL;
	particle* p = set->particles;
	int i;
#ifdef DEBUG
	printf("mover1d_postprocessing: start\n");
#endif

	CUDA_CALL(cudaMemcpy(&d_zpool, &(mover->zpool), sizeof d_zpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&d_vxpool, &(mover->vxpool), sizeof d_vxpool, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&d_vypool, &(mover->vypool), sizeof d_vypool, cudaMemcpyDeviceToHost));

	CUDA_CALL(cudaMalloc(&d_rvx, sizeof(floatdig)*PARTICLEPOOL));
	CUDA_CALL(cudaMalloc(&d_rvy, sizeof(floatdig)*PARTICLEPOOL));
	for (i = 0; i < nbunches; i++) {
		CURAND_CALL(curandGenerateNormal(set->cuda_r, d_rvx, PARTICLEPOOL, 0, sigma));
		CURAND_CALL(curandGenerateNormal(set->cuda_r, d_rvy, PARTICLEPOOL, 0, sigma));
		copypoolHostToDevice(NULL, NULL, NULL, d_zpool, p, i*PARTICLEPOOL);
		do_postproceesing<<<NUMOFTHREADBLOCKS,NUMOFTHREADSPERBLOCK>>>(d_vxpool, d_vypool, d_zpool, d_rvx, d_rvy, minz, maxz);
		copypoolDeviceToHost(d_vxpool, d_vypool, NULL, NULL, NULL, d_zpool, p, i*PARTICLEPOOL);
	}

	CUDA_CALL(cudaFree(d_rvx));
	CUDA_CALL(cudaFree(d_rvy));

#ifdef DEBUG
	printf("mover1d_postprocessing: done\n");
#endif

}

void mover1d_dump_E_field_into_file(mover1d* mover, char* path, int fid, floatdig minz, floatdig maxz, floatdig dz, floatdig t)
{
	floatdig z;
	char fname[MAX_BUF_SIZE];
	FILE* to;
	floatdig e;
	floatdig* d_zpool;
	floatdig* d_E;
	floatdig* d_B;
	sprintf(fname, "%s/%s_%d.dat", HOMEDIR, EFIELDDUMPFILEZ, fid);
	to = fopen(fname, "w");
	CUDA_CALL(cudaMemcpy(&d_E, &(mover->E), sizeof d_E, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&d_B, &(mover->B), sizeof d_B, cudaMemcpyDeviceToHost));
	CUDA_CALL(cudaMemcpy(&d_zpool, &(mover->zpool), sizeof d_zpool, cudaMemcpyDeviceToHost));
	for (z = minz; z <= maxz; z += dz) 
	{
		CUDA_CALL(cudaMemcpy(d_zpool, &z, sizeof z, cudaMemcpyHostToDevice));
		update_external_fields<<<1,1>>>(d_E, d_B, d_zpool, t);
		CUDA_CALL(cudaMemcpy(&e, d_E, sizeof e, cudaMemcpyDeviceToHost));
		fprintf(to, "%e\n", e);
	}
	fclose(to);

}
